package com.haiyou.common.util.limiter;

import java.util.concurrent.TimeUnit;

import com.github.benmanes.caffeine.cache.Cache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.google.common.util.concurrent.RateLimiter;

import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

/**
 * 
 * 
 * @Description: 基于 key 的限流器
 * @author xingyuan
 * @date 2022年11月2日 上午10:58:11
 */
@Slf4j
public class KeyRateLimiter {

	/**
	 * 缓存
	 */
	private Cache<String, RateLimiter> cache;

	/**
	 * 每秒频率
	 */
	@Getter
	private double permitsPerSecond;
	/**
	 * key 过期时间
	 */
	@Getter
	private int expireAfterAccess;

	private KeyRateLimiter(int maximumSize,int expireAfterAccess,double permitsPerSecond) {
		
		this.permitsPerSecond = permitsPerSecond;
		
		this.expireAfterAccess = expireAfterAccess;
		
		this.cache = Caffeine.newBuilder()
				//设置最大条目
				.maximumSize(maximumSize)
				//设置读写过去时间 单位s
				.expireAfterAccess(expireAfterAccess, TimeUnit.SECONDS)
				
				.build();
	}

	/**
	 * 尝试获取令牌
	 * @param key
	 * @param permits
	 * @param timeout
	 * @param unit
	 * @return
	 */
	public boolean tryAcquire(String key,int permits, long timeout, TimeUnit unit) {
		return rateLimiter(key).tryAcquire(permits, timeout, unit);
	}
	
	
	/**
	 * 基于key 限流
	 * @param key
	 * @return
	 */
	public boolean tryAcquire(String key) {
		return tryAcquire(key,1,0,TimeUnit.MICROSECONDS);
	}
	

	/**
	 * 获取 key 对应限流器
	 * 
	 * @param key
	 * @return
	 */
	private RateLimiter rateLimiter(String key) {
		return cache.get(key, k -> {
			
			RateLimiter limiter = RateLimiter.create(permitsPerSecond);
			
			return limiter;
		});
	}

	/**
	 * 初始化创建
	 * 
	 * @param maximumSize       最大key条目数量
	 * @param expireAfterAccess key 读写缓存时间 秒
	 * @param permitsPerSecond  每秒频率
	 */
	public static KeyRateLimiter create(int maximumSize, int expireAfterAccess, double permitsPerSecond) {
		return new KeyRateLimiter(maximumSize, expireAfterAccess, permitsPerSecond);
	}
	
}
